import unittest
from dev.entity.stock import Stock
import pandas as pd

import sys
import os
from pathlib import Path
curPath = os.path.abspath(os.path.dirname(__file__))
rootPath = os.path.split(curPath)[0]
sys.path.append(rootPath)

stock_code_path = os.path.join(Path(rootPath),"entity/temp/","stock_code.csv")

def get_stocks_name(stock_codes_path=stock_code_path):
    res_df = pd.read_csv(stock_codes_path)
    codes = list(res_df['代码'].values)
    names = list(res_df['名称'].values)
    code_name_dict = {}
    for i in range(len(codes)):
        code_name_dict[codes[i]] = names[i]
    return code_name_dict

code_name_dict = get_stocks_name(stock_code_path)

class StockTest(unittest.TestCase):

    #测试Stock类初始化
    def test_Stock_init(self):
        # code = 'sh601728'  #中国电信
        code = 'sz000722'  #湖南发展
        codeName = code_name_dict[code]
        saveDir = os.path.join(rootPath,"entity/temp/",code)
        stock = Stock(code=code,codeName=codeName,saveDir=saveDir) #初始化股票Stock对象
        stock.plot_K_line()

    # 测试Stock类设置分时K线对象
    def test_Stock_set_minute_K_line(self):
        # code = 'sh601728'  # 中国电信
        code = 'sz000722'  # 湖南发展
        codeName = code_name_dict[code]
        saveDir = os.path.join(rootPath, "entity/temp/", code)
        stock = Stock(code=code, codeName=codeName, saveDir=saveDir)  # 初始化股票Stock对象
        stock.set_minute_K_line(period="1",indices_params={})
        stock.plot_K_line()

    # 测试Stock类设置日K线对象
    def test_Stock_set_daily_K_line(self):
        # code = 'sh601728'  # 中国电信
        code = 'sz000722'  # 湖南发展
        codeName = code_name_dict[code]
        saveDir = os.path.join(rootPath, "entity/temp/", code)
        stock = Stock(code=code, codeName=codeName, saveDir=saveDir)  # 初始化股票Stock对象

        start_date, end_date = "20221101","20230217"
        stock.set_daily_K_line(start_date=start_date, end_date=end_date, indices_params={})
        stock.plot_K_line()

    # 测试Stock类设置周K线对象
    def test_Stock_set_week_K_line(self):
        code = 'sh601728'  # 中国电信
        # code = 'sz000722'  # 湖南发展
        codeName = code_name_dict[code]
        saveDir = os.path.join(rootPath, "entity/temp/", code)
        stock = Stock(code=code, codeName=codeName, saveDir=saveDir)  # 初始化股票Stock对象

        start_date, end_date = "20200101", "20230220"
        stock.set_week_K_line(start_date=start_date, end_date=end_date, indices_params={})
        stock.plot_K_line()

    # 测试Stock类设置月K线对象（不适合绘制MACD（12，26，9）和RSI（6，12，24），原因是时间窗口不够）
    def test_Stock_set_month_K_line(self):
        # code = 'sh601728'  # 中国电信
        code = 'sz000722'  # 湖南发展
        codeName = code_name_dict[code]
        saveDir = os.path.join(rootPath, "entity/temp/", code)
        # 标签使用标志
        indices_dict = {
            "MACD": False,
            "RSI": False,
            "Bollinger": True,
            "obv": True
        }
        stock = Stock(code=code, codeName=codeName, saveDir=saveDir)  # 初始化股票Stock对象

        start_date, end_date = "20220101", "20230217"
        stock.set_month_K_line(start_date=start_date, end_date=end_date, indices_params={},indices_useflag_dict=indices_dict)
        stock.plot_K_line()


